{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OchreLang.Syntax
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Two level syntax - pure values and effectful expressions.
--
--------------------------------------------------------------------------------

module OchreLang.Syntax
  (
    Orchestra(..)
  , InstrumentSource(..)
  , Instrument(..)
  , MethodDecl(..)
  , MethodDefn(..)
  , Expr(..)
  , Value(..)
  
  , effect

  -- Traversals
  , mapM_InstrumentSource
  , mapM_Instrument
  , foreachM_Instrument

  -- Pretty
  , ppOrchestra
  , ppInstrument
  , ppExpr
  , ppValue

  , InstrumentDict
  , buildInstrumentDict
  , findOrchestra


  , MethodDict
  , buildMethodDict
  , findMethodDefn

  , FreeVars
  , freevars

  , PrimIds
  , primIds

  ) where

import qualified OchreLang.OrchSpec as OS
import OchreLang.SyntaxBase
import OchreLang.Utils

import Text.PrettyExprHPJ                       -- package: pretty-expr-hpj
import Text.PrettyExprHPJ.Haskell98

import Text.PrettyPrint.HughesPJ                -- package: pretty

import Control.Applicative
import qualified Data.Map as Map
import Data.Set ( (\\) )
import qualified Data.Set as Set



-- Note - for super-inlining we need to consider all instruments 
-- in a module not just the ones included in an orchestra...
--

data Orchestra = Orchestra
    { orch_instrs       :: [InstrumentSource]
    }
  deriving (Eq,Show)



-- | Alternative view of Instruemnt compounding information 
-- from the Orchestra declaration.
--
data InstrumentSource = InstrumentSource 
    { instr_src_num     :: Int
    , instr_src_inits   :: [OS.InitDecl]
    , instr_src_instr   :: Instrument
    } 
  deriving (Eq,Show)

-- | Instrument name is residual from top-level syntax, it should 
-- help debugging...
--
data Instrument = Instrument
    { instr_name        :: Ident
    , instr_superclass  :: Ident
    , instr_main        :: Expr
    , instr_methods     :: [MethodDecl]
    , instr_module      :: ModuleName
    }
  deriving (Eq,Show)

data MethodDecl = MethodDecl Ident MethodDefn
  deriving (Eq,Show)

data MethodDefn = MethodDefn [VarId] Expr
  deriving (Eq,Show)


infixl 1 :>>

-- | Prim is prim-application.
--
-- LetFun uses a VarId rather than Ident so the name can be alpha-renamed
--
data Expr = Return      Value                         -- lift a pure Value (needed?)
          | OpCall      Ident     [Value]
          | SuperCall   Ident     [Value]
          | FunApp      VarId     [Value]
          | LetValue    VarId     Value     Expr          -- float, int, bool, ...
          | LetSig1     VarId     Expr      Expr          -- ISig, KSig, ASig, ...
          | LetSig2     VarId     VarId     Expr    Expr
          | LetSig3     VarId     VarId     VarId   Expr    Expr
          | LetSig4     VarId     VarId     VarId   VarId   Expr    Expr
          | LetTuple    [VarId]   Expr      Expr          -- eliminated after expand_main
          | LetFun      VarId     [VarId]   Expr    Expr
          | Ifte        Value     Expr      Expr
          | Expr        :>>       Expr                           -- sequence
  deriving (Eq,Show)

-- TODO (?) - (| LetSigTup [VarId] Expr Expr |) would be more general than 
-- LetSig2, LetSig3, LetSig4...


-- | Separation of Values and Expressions.
--
-- Bool helps const folding even though Csound doesn\'t have one.
--
-- Warning - Pfield is anti-modular and should only be introduced by 
-- init-expressions.
--
data Value = Unit
           | Bool       Bool
           | Int        Int
           | Float      Decimal
           | String     String
           | Var        VarId
           | Pfield     Int
           | Tuple      [Value]
           | UnaryV     UnaryOp   Value
           | BinV       BinOp     Value     Value
           | RelV       RelOp     Value     Value
           | CondV      Value     Value     Value
           | FunCallV   Ident     [Value]
  deriving (Eq,Show)


-- 
-- Note - Csound if expressions must have a binary relational 
-- expression for the cond expression, the following
-- will not compile:
-- 
-- > if 0 then ...
--
-- The syntax allows this so we must do something in the 
-- translation...
--



-- | Note the type of an Opcode call (Prim) is its answer type
-- (i.e. the application of Prim to its arguments) not the type 
-- of the opcode function.
--
-- C.f. Haskell
--
-- > :t ( 1 + 2)  == Int  (Prim App)
-- 
-- > :t (+)       == Int -> Int -> Int (Prim Def)
--
--



-- | Sequencing is expected to be effectful in the left...
-- 
effect :: Expr -> Bool
effect (Return {})              = False
effect (OpCall {})              = True
effect (SuperCall {})           = True          -- Assume True should have been inlined...
effect (FunApp {})              = True
effect (LetValue _ _ e1)        = effect e1
effect (LetSig1 _ e1 e2)        = effect e1 || effect e2
effect (LetSig2 _ _ e1 e2)      = effect e1 || effect e2
effect (LetSig3 _ _ _ e1 e2)    = effect e1 || effect e2
effect (LetSig4 _ _ _ _ e1 e2)  = effect e1 || effect e2
effect (LetTuple _ e1 e2)       = effect e1 || effect e2
effect (LetFun _ _ e1 e2)       = effect e1 || effect e2
effect (Ifte _ e1 e2)           = effect e1 || effect e2
effect (e1 :>> e2)              = effect e1 || effect e2



--------------------------------------------------------------------------------
-- Traversal


mapM_InstrumentSource :: Monad m 
                      => (InstrumentSource -> m InstrumentSource) 
                      -> Orchestra 
                      -> m Orchestra
mapM_InstrumentSource mf orch@(Orchestra { orch_instrs = srcs }) =
   do { ans <- mapM mf srcs
      ; return $ orch { orch_instrs = ans }
      }


mapM_Instrument :: Monad m 
                => (Instrument -> m Instrument) 
                -> Orchestra 
                -> m Orchestra
mapM_Instrument mf = mapM_InstrumentSource inner
  where
    inner src@(InstrumentSource { instr_src_instr = inst }) = do 
        { ans <- mf inst
        ; return $ src { instr_src_instr = ans }
        }

foreachM_Instrument :: Monad m
                    => (Instrument -> m ()) -> Orchestra -> m ()
foreachM_Instrument mf (Orchestra { orch_instrs = srcs }) = mapM_ inner srcs
  where
    inner (InstrumentSource { instr_src_instr = inst }) = mf inst


--------------------------------------------------------------------------------
-- pretty print

ppOrchestra :: Orchestra -> Doc
ppOrchestra (Orchestra instrs) = 
    text "orch" $+$ (vsemiBraces $ map ppInstrumentSource instrs)


ppInstrumentSource :: InstrumentSource -> Doc
ppInstrumentSource  (InstrumentSource { instr_src_num     = i 
                                      , instr_src_inits   = idecls
                                      , instr_src_instr   = instr } ) = 
       text "instr_num" <+> char '=' <+> int i
   $+$ text "inits" <+> char '=' <+> (vsemiBraces $ map OS.ppInitDecl idecls)
   $+$ ppInstrument instr


ppInstrument :: Instrument -> Doc
ppInstrument (Instrument { instr_name       = name
                         , instr_superclass = super
                         , instr_main       = mexpr
                         , instr_methods    = mdecls }) = 
    text "instr" <+> ppIdent name <+> text "extends" <+> ppIdent super
      $+$ nest 2 (vsemiBraces $ ds ++ [d1])
  where
    d1 = ppMethodDecl $ MethodDecl (Ident "main") (MethodDefn [] mexpr)
    ds = map ppMethodDecl mdecls

ppMethodDecl :: MethodDecl -> Doc
ppMethodDecl (MethodDecl name defn) = 
    ppIdent name <+> char '=' <+> ppMethodDefn defn


ppMethodDefn :: MethodDefn -> Doc
ppMethodDefn  (MethodDefn args body)
    | null args = ppExpr body
    | otherwise = let formals = hsep $ map ppVarId args
                  in char '\\' <> formals <+> text "->" <+> ppExpr body


ppExpr :: Expr -> Doc
ppExpr = unparse . exprDoc



exprDoc :: Expr -> DocE
exprDoc (Return v)                  = valueDoc v
exprDoc (OpCall vid vals)           = 
    Atom $ ppIdent vid <+> commaSep (map ppValue vals)

exprDoc (SuperCall v vals)          = 
    Atom $ text "SUPER:" <> ppIdent v <+> commaSep (map ppValue vals)

exprDoc (FunApp v vals)             = 
    Atom $ ppVarId v <+> commaSep (map ppValue vals)


exprDoc (LetValue vid v1 e1)        = 
    let pre = text "let" <+> ppVarId vid
    in Atom $ letBody pre (ppValue v1) (ppExpr e1)

exprDoc (LetSig1 vid e1 e2)         = 
    let pre = text "letsig" <+> ppVarId vid
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (LetSig2 v1 v2 e1 e2)       = 
    let pre = text "letsig2" <+> tupled (map ppVarId [v1,v2])
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (LetSig3 v1 v2 v3 e1 e2)    = 
    let pre = text "letsig3" <+> tupled (map ppVarId [v1,v2,v3])
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (LetSig4 v1 v2 v3 v4 e1 e2) = 
    let pre = text "letsig4" <+> tupled (map ppVarId  [v1,v2,v3,v4])
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (LetTuple vs e1 e2)         = 
    let pre = text "lettuple" <+> tupled (map ppVarId vs)
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (LetFun v args e1 e2)       = 
    let pre = text "letfun" <+> ppVarId v <+> hsep (map ppVarId args) 
    in Atom $ letBody pre (ppExpr e1) (ppExpr e2)

exprDoc (Ifte v1 e1 e2)             = 
    Atom $ text "if" <+> ppValue v1
                     <+> nest 2 (    text "then" <+> ppExpr e1
                                 $+$ text "else" <+> ppExpr e2)

exprDoc (e1 :>> e2)                 = monadicSeqB (exprDoc e1) (exprDoc e2)  


ppValue :: Value -> Doc
ppValue = unparse . valueDoc

valueDoc :: Value -> DocE
valueDoc (Unit)                     = literal "()"
valueDoc (Bool True)                = Atom $ text "true" 
valueDoc (Bool False)               = Atom $ text "false"
valueDoc (Int i)                    = Atom $ int i
valueDoc (Float d)                  = 
    Atom $ text $ truncatedDouble $ realToFrac d

valueDoc (String ss)                = Atom $ doubleQuotes $ text ss
valueDoc (Var v)                    = literal $ varName v 
valueDoc (Pfield i)                 = literal $ 'p' : show i
valueDoc (Tuple vs)                 = Atom $ tupled $ map ppValue vs 
valueDoc (UnaryV op v)              = unaryDocE op (valueDoc v)
valueDoc (BinV op v1 v2)            = binDocE op (valueDoc v1) (valueDoc v2)
valueDoc (RelV relop v1 v2)         = 
    relDocE relop (valueDoc v1) (valueDoc v2)

valueDoc (CondV vc v1 v2)           = 
    Atom $ text "cond" <+> ppValue vc <+> char '?' <+> ppValue v1
                                      <+> char ':' <+> ppValue v2

valueDoc (FunCallV v es)            = 
    funAppB (literal $ getIdent v) (map valueDoc es)


letBody :: Doc -> Doc -> Doc -> Doc
letBody d e body = 
    d <+> char '=' <+> e <+> text "in" $+$ nest 2 body


--------------------------------------------------------------------------------
-- 

type InstrumentDict = Map.Map Ident Instrument


buildInstrumentDict :: [Instrument] -> InstrumentDict
buildInstrumentDict = foldr f1 Map.empty
  where
    f1 ins   ac         = Map.insert (instr_name ins) ins ac


findInstrument :: Ident -> InstrumentDict -> Maybe Instrument
findInstrument = Map.lookup


-- | Note - we don\'t want this in the compiler monad as it 
-- drags in an unwanted dependency.
-- 
findOrchestra :: [OS.InstrInstance] 
                -> InstrumentDict 
                -> Either String Orchestra
findOrchestra xs dict = Orchestra <$> foldr fn (Right []) xs
  where
    fn _  (Left err)  = Left err
    fn ii (Right ac)  = let vid = OS.instance_instr_name ii in 
        case findInstrument vid dict of
          Nothing -> Left $ "Missing instrument: " ++ getIdent vid
          Just ins -> let ans = InstrumentSource 
                                  { instr_src_num   = OS.instance_number ii
                                  , instr_src_inits = OS.instance_init_decls ii
                                  , instr_src_instr = ins
                                  }
                      in Right $ ans:ac

type MethodDict     = Map.Map ScopedMethodName MethodDecl

data ScopedMethodName = ScopedMethodName 
    { scoped_instr_name         :: !String      -- aka Class
    , scoped_method_name        :: !String
    }
  deriving (Eq,Ord,Show)

allMethods :: Instrument -> [MethodDecl]
allMethods (Instrument { instr_main = main_expr, instr_methods = ms}) =
    let main_decl = MethodDecl (Ident "main") (MethodDefn []  main_expr)
    in main_decl : ms

makeScopedMethodName :: Ident -> MethodDecl -> ScopedMethodName
makeScopedMethodName ins_name (MethodDecl vid _) = 
    ScopedMethodName { scoped_instr_name  = getIdent ins_name
                     , scoped_method_name = getIdent vid
                     }


-- | At some time will need to consider modules and scoped 
-- names...
--
buildMethodDict :: InstrumentDict -> MethodDict
buildMethodDict = foldr f1 Map.empty . Map.elems
  where
    f1 ins ac           = 
        let ms = allMethods ins in foldr (f2 (instr_name ins)) ac ms

    f2 ins_name decl ac = 
        let mname = makeScopedMethodName ins_name decl in Map.insert mname decl ac


findMethodDefn :: Ident -> Ident -> MethodDict -> Maybe MethodDefn
findMethodDefn ins_name method_name dict = 
    fmap extract $ Map.lookup full_name dict
  where
    extract (MethodDecl _ body) = body
    full_name = ScopedMethodName { scoped_instr_name  = getIdent ins_name
                                 , scoped_method_name = getIdent method_name }

--------------------------------------------------------------------------------
-- Free vars

type FreeVars = Set.Set VarId

freevars :: Expr -> FreeVars
freevars (Return v)                     = freevarsV v

freevars (OpCall _ xs)                  = Set.unions $ map freevarsV xs

freevars (SuperCall _ xs)               = Set.unions $ map freevarsV xs

freevars (FunApp v xs)                  = 
    Set.insert v $ Set.unions $ map freevarsV xs

freevars (LetValue v val e2)            =
    Set.union (freevarsV val) (Set.delete v $ freevars e2)

freevars (LetSig1 v e1 e2)              =
    Set.union (freevars e1) (Set.delete v $ freevars e2)

freevars (LetSig2 v1 v2 e1 e2)          =
    let lhs = freevars e1
        rhs = Set.delete v2 $ Set.delete v1 $ freevars e2
    in Set.union lhs rhs
                             
freevars (LetSig3 v1 v2 v3 e1 e2)       = 
    let lhs = freevars e1
        rhs = Set.delete v3 $ Set.delete v2 $ Set.delete v1 $ freevars e2
    in Set.union lhs rhs

freevars (LetSig4 v1 v2 v3 v4 e1 e2)    = 
    let lhs = freevars e1
        rhs = Set.delete v4 $ Set.delete v3 $ Set.delete v2 $ Set.delete v1 $ 
                freevars e2
    in Set.union lhs rhs

freevars (LetTuple vs e1 e2)            = 
    let lhs = freevars e1
        rhs = foldr Set.delete (freevars e2) vs
    in Set.union lhs rhs


-- This definition isn\'t really correct 
freevars (LetFun v args e1 e2)          = 
    let zs  = freevars e1 \\ Set.fromList args
        rhs = freevars e2
    in Set.union zs rhs \\ Set.singleton v

freevars (Ifte v1 e1 e2)                = 
    Set.unions [freevarsV v1, freevars e1, freevars e2]

freevars (e1 :>> e2)                    = Set.union (freevars e1) (freevars e2)


freevarsV :: Value -> FreeVars
freevarsV (Unit)                 = Set.empty
freevarsV (Bool {})              = Set.empty
freevarsV (Int {})               = Set.empty
freevarsV (Float {})             = Set.empty
freevarsV (String {})            = Set.empty
freevarsV (Var v)                = Set.singleton v
freevarsV (Pfield {})            = Set.empty
freevarsV (Tuple vs)             = Set.unions $ map freevarsV vs
freevarsV (UnaryV _ v)           = freevarsV v
freevarsV (BinV _ v1 v2)         = freevarsV v1 `Set.union` freevarsV v2
freevarsV (RelV _ v1 v2)         = freevarsV v1 `Set.union` freevarsV v2
freevarsV (CondV v1 vt vf)       = Set.unions $ map freevarsV [v1,vt,vf]
freevarsV (FunCallV _ xs)        = Set.unions $ map freevarsV xs


-- freevarsV _                      = Set.empty


--------------------------------------------------------------------------------

type PrimIds = Set.Set Ident


primIds :: Expr -> PrimIds
primIds (Return v)                  = primIdsV v
primIds (OpCall s xs)               = let s1 = Set.singleton s
                                          ss = map primIdsV xs 
                                      in Set.unions $ s1:ss
primIds (SuperCall _ xs)            = Set.unions $ map primIdsV xs
primIds (FunApp _ xs)               = Set.unions $ map primIdsV xs
primIds (LetValue _ val e2)         = primIdsV val `Set.union` primIds e2
primIds (LetSig1 _ e1 e2)           = primIds e1 `Set.union` primIds e2
primIds (LetSig2 _ _ e1 e2)         = primIds e1 `Set.union` primIds e2
primIds (LetSig3 _ _ _ e1 e2)       = primIds e1 `Set.union` primIds e2
primIds (LetSig4 _ _ _ _ e1 e2)     = primIds e1 `Set.union` primIds e2
primIds (LetTuple _ e1 e2)          = primIds e1 `Set.union` primIds e2
primIds (LetFun _ _ e1 e2)          = primIds e1 `Set.union` primIds e2
primIds (Ifte v1 e1 e2)             = 
    Set.unions [primIdsV v1, primIds e1, primIds e2]

primIds (e1 :>> e2)                 = Set.union (primIds e1) (primIds e2)


primIdsV :: Value -> PrimIds
primIdsV (Var {})               = Set.empty
primIdsV (Tuple vs)             = Set.unions $ map primIdsV vs
primIdsV (UnaryV _ v)           = primIdsV v
primIdsV (BinV _ v1 v2)         = primIdsV v1 `Set.union` primIdsV v2
primIdsV (RelV _ v1 v2)         = primIdsV v1 `Set.union` primIdsV v2
primIdsV (CondV v1 vt vf)       = Set.unions $ map primIdsV [v1,vt,vf]
primIdsV (FunCallV s xs)        = let s1 = Set.singleton s
                                      ss = map primIdsV xs 
                                  in Set.unions $ s1:ss
primIdsV _                      = Set.empty


